import jieba
import pickle

from data_source import FileSource
from config import pos_path, neg_path, data_path
from dict_model.WordFilter import filter_stop
from naive_bayes_model.bayes import NaiveBayes


def sent2word(sentence):
    word_gen = jieba.cut(sentence, cut_all=False)
    word_list = []
    for word in word_gen:
        word_list.append(word)
    word_list = filter_stop(word_list)
    return word_list


class Sentiment:
    def __init__(self):
        self.classifier = NaiveBayes()
        self.pos_source = FileSource()
        self.neg_source = FileSource()

    def train(self):
        self.pos_source.open(pos_path)
        self.neg_source.open(neg_path)
        data = []
        for line in self.pos_source:
            words = sent2word(line)
            data.append([words, 1])
        for line in self.neg_source:
            words = sent2word(line)
            data.append([words, -1])
        self.classifier.train(data)

    def predict(self, sentence):
        words = sent2word(sentence)
        result, prob = self.classifier.predict(words)
        if prob < 0.65:
            result = 'neg'
            prob = 1 - prob
        else:
            result = 'pos'
        return result, prob

    def save(self, **kwargs):
        filename = kwargs.get('filename')
        if filename:
            filename = data_path(filename)
        else:
            filename = data_path('bayes.model')
        fw = open(filename, 'wb')
        pickle.dump(self.classifier, fw)

    def load(self, **kwargs):
        filename = kwargs.get('filename')
        if filename:
            filename = data_path(filename)
        else:
            filename = data_path('bayes.model')
        fr = open(filename, 'rb')
        self.classifier = pickle.load(fr)

    def test(self):
        cases = [
            '大家好才是真的好',
            '这是一本好书',
            '这本书就是垃圾',
            '这种玩意还是扔掉好',
            '非常非常喜欢'
        ]
        for case in cases:
            result, prob = self.predict(case)
            print(case, result)


if __name__ == '__main__':
    s = Sentiment()
    # s.train()
    # s.save()
    s.load()
    s.test()
